# coding:utf-8

import os
import torch
from torch import nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision.transforms import transforms
from Unet import *
from data_pk import *

use_gpu = torch.cuda.is_available()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  # 定义训练使用的设备

###########新建数据集###############
inputlist = getFilepathInDir("../oct_img/11")  # ("../oct_img/input")
saveinavi = "../oct_img/input.avi"
saveoutavie = "../oct_img/output.avi"

reqs = (512, 512)

dataset = ImgListDataSet(inputlist, reqsize=reqs)
totensor = transforms.ToTensor()

fourcc = cv.VideoWriter_fourcc(*'XVID')
writer1 = cv.VideoWriter(saveinavi, fourcc, 24.0, reqs)
writer2 = cv.VideoWriter(saveoutavie, fourcc, 24.0, reqs)

#########新建模型##############
# model = UNet(in_channel=1, out_channel=1)
# model = LessUNet(in_channel=1, out_channel=2)
# model.add_module("non_linear_layer", MySigmoid())
model = MixUNet1()
model: torch.nn.Module
# model.half()  # 使用半精度推理 节省计算量
model.to(device)

#############加载模型参数#############
# savemoduleparamName = "../moduleparam/enhance_unet_sigmoid.pth"
# savemoduleparamName = "../moduleparam/enhance_lessunet2.pth"
# savemoduleparamName = "../moduleparam/enhance_mixlessunet1.pth"
# savemoduleparamName = "../moduleparam/enhance_mixlessunet2.pth"
# savemoduleparamName = "../moduleparam/enhance_mixlessunet3.pth"
savemoduleparamName = "../moduleparam/enhance_mixlessunet4.pth"
if os.path.exists(savemoduleparamName):
    model.load_state_dict(torch.load(savemoduleparamName))
else:
    newdir, _ = os.path.split(savemoduleparamName)
    os.makedirs(newdir, exist_ok=True)

############# 开始进行推理 ####################
model.eval()  # 确保为评估模式  不再做dropout  normalization
with torch.no_grad():  # with no_grad()节省推理时内存
    for i, img in enumerate(dataset):
        img: np.ndarray
        if i >= len(inputlist):
            break
        print("图片 %d  - %s" % (i, inputlist[i]))
        cv.imshow("input", img)
        cv.waitKey(1)
        input = totensor(img)
        input: torch.Tensor
        input = torch.unsqueeze(input, 0).to(device)
        # input = input.half()
        output = model(input)
        img_output = (output.detach()[0] * 255).cpu().numpy()
        img_output = np.transpose(img_output, (1, 2, 0)).astype(np.uint8)

        cv.imshow("output", img_output)
        cv.waitKey(1)

        img = cv.cvtColor(img, cv.COLOR_GRAY2BGR)
        img_output = cv.cvtColor(img_output, cv.COLOR_GRAY2BGR)
        writer1.write(img)
        writer2.write(img_output)
